import cv2
from torch.fx.experimental.unification.utils import xfail

img = cv2.imread("../images/02.png")

# 1. 灰度化
img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

# 2. 二值化
_, img_binary = cv2.threshold(
    img_gray,
    127,
    255,
    cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU
)

# 3. 高斯滤波
img_gaussian = cv2.GaussianBlur(
    img_binary,
    (5, 5),
    3
)

# 4. 寻找轮廓
contours, _ = cv2.findContours(
    img_gaussian,
    mode=cv2.RETR_LIST,  # 查找轮廓的方式
    method=cv2.CHAIN_APPROX_SIMPLE  # 轮廓近似方法
)
print(f"寻找到的轮廓的个数为：{len(contours)}")
for i, c in enumerate(contours):
    print(f"第{i + 1}个轮廓的边界点的个位数为{len(c)}")

# 5. 循环遍历轮廓列表，做多边形逼近【把一个复杂的轮廓用较少的顶点近似表示】
for cnt in contours:
    # 5.1 计算轮廓的周长
    perimeter = cv2.arcLength(cnt, True)
    # 5.2 根据周长确定 epsilon (原轮廓与近似多边形的最大距离) 精度，
    approx_pts = cv2.approxPolyDP(cnt, epsilon=perimeter * 0.04, closed=True)
    # 5.3 根据逼近后的顶点，绘制逼近后的轮廓
    cv2.drawContours(img, [approx_pts], -1, (0, 0, 255), 2)
    # 5.4 对比原有轮廓
    cv2.drawContours(img, [cnt], -1, (255, 0, 0), 1)
    # 5.5 判断逼近后的轮廓的顶点个数，确定形状
    shape = "None"
    if len(approx_pts) == 3:
        shape = "triangle"
    elif len(approx_pts) == 4:
        # 怎么进一步确定 正方形 还是 矩形
        # 根据长宽度确定，正方形的长宽比为1，当然我们允许有 5% 的误差
        x, y, w, h = cv2.boundingRect(approx_pts)
        if 0.95 <= w / h <= 1.05:
            shape = "square"
        else:
            shape = "rectangle"
    elif len(approx_pts) == 5:
        shape = "pentagon"
    elif len(approx_pts) >= 8:  # 顶点数较多，可能是圆
        # 进一步验证是否为圆
        area = cv2.contourArea(cnt)
        perimeter = cv2.arcLength(cnt, True)
        # 圆的特性：perimeter²/(4*π*area) ≈ 1
        circularity = perimeter * perimeter / (4 * 3.14159 * area)
        if 0.7 <= circularity <= 1.3:  # 允许一定误差
            shape = "circle"
        else:
            shape = "unknown"
    elif len(approx_pts) == 2:
        shape = "line"
    print(len(approx_pts), shape)
    # 5.6 将形状文字 标注到图形上
    # cv2.putText(img, text=shape, org=(10,10), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5, color=(0,255,255))
    # 上面这种写法不报错，但是文字的位置不对，为了将文字写在识别到的形状上，可以借助cv2.moments()
    M = cv2.moments(cnt)
    # 质心的 x 坐标为 cx = m10 / m00，质心的 y 坐标为 cy = m01 / m00
    cx = int(M["m10"] / M["m00"])
    cy = int(M["m01"] / M["m00"])
    cv2.putText(img, text=shape, org=(cx, cy), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=0.5, color=(0, 0, 0))

cv2.imshow("image", img)
cv2.waitKey(0)
